import websockets
import asyncio
import socket

from main.asr import ASR
from main.tts import TTS
from main.llm import LLM
from main.vad import VAD
from main.handle import ConectionHandle

class WebsocketServer:
    
    def __init__(self):
        self.active_connections = set()
        self._vad = VAD()
        self._asr = ASR()
        self._tts = TTS()
        self._llm = LLM()
    
    def get_local_ip(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except Exception:
            return None

    async def start(self):
        host = "127.0.0.1" #self.get_local_ip()
        port = 8000
        print(host, port)
        
        ws = await websockets.serve(self._handle_Coonection, host=host, port=port)
        await ws.wait_closed()
    
    async def _handle_Coonection(self, websocket):
        print("链接成功")
        handler = ConectionHandle(self._vad, self._asr, self._llm, self._tts)
        self.active_connections.add(handler)
        try:
            await handler.handle_connection(websocket=websocket)
        finally:
            self.active_connections.discard(handler)

if __name__ == "__main__":
    ws = WebsocketServer()
    asyncio.run(ws.start())
    